/* Copyright (c) 2011-2013 ISP RAS, Moscow. For more information please write to

      mermakov@ispras.ru
      svartanov@ispras.ru

   Licensed under the Apache License, Version 2.0 (the "License");
   you may not use this file except in compliance with the License.
   You may obtain a copy of the License at

       http://www.apache.org/licenses/LICENSE-2.0

   Unless required by applicable law or agreed to in writing, software
   distributed under the License is distributed on an "AS IS" BASIS,
   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
   See the License for the specific language governing permissions and
   limitations under the License.
*/

package ru.ispras.coffeemachine.instrument;

import java.io.IOException;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.HashSet;
import java.util.List;
import java.util.Set;

import org.apache.bcel.Constants;
import org.apache.bcel.Repository;
import org.apache.bcel.classfile.JavaClass;
import org.apache.bcel.classfile.LineNumberTable;
import org.apache.bcel.classfile.Method;
import org.apache.bcel.generic.BranchInstruction;
import org.apache.bcel.generic.ClassGen;
import org.apache.bcel.generic.CodeExceptionGen;
import org.apache.bcel.generic.ConstantPoolGen;
import org.apache.bcel.generic.InstructionHandle;
import org.apache.bcel.generic.InstructionList;
import org.apache.bcel.generic.MethodGen;
import org.apache.bcel.generic.PUSH;
import org.apache.bcel.generic.TargetLostException;

import ru.ispras.coffeemachine.core.CoffeeMachine;
import ru.ispras.coffeemachine.instrument.util.Utility;
import ru.ispras.coffeemachine.util.Logger;
import ru.ispras.coffeemachine.util.Options;

/**
 * Mokka, class file instrumentation
 *
 * @author Sergey Vartanov (me@enzet.ru)
 * @author Mikhail Ermakov (mermakov@ispras.ru)
 * @version 1.2.0
 * @since 2011-09-12
 */
public class Mokka
{
	private final JavaClass originCode;
	private final ClassGen modifiedCode;

	private final String className;

	static int pc = 0;

	/**
	 * List of methods which must be instrumented. If <tt>null</tt> all methods
	 * will be instrumented
	 */
	private Set<String> methods;

	/**
	 * Extended instruction list.
	 *
	 * Normally during instrumentation for each instruction <b>I</b> on position <b>X</b>
	 * Coffee Machine generates a set of instructions that replaces <b>I</b> moving it further on position <b>Y</b>.
	 * BCEL automatically modifies all jump instructions that point to position <b>X</b> in current method
	 * replacing target <b>X</b> with target <b>Y</b>. This is, however, wrong for Coffee Machine purposes and
	 * is rolled back after full method instrumentation.
	 *
	 * For several cases jump instruction targets must not point to the beginning of inserted list.
	 * This is most prominently used for extra instruction sets inserted on position <code>0</code>, such
	 * as method entry notifications. Therefore, an extended instruction list is used, which has
	 * a custom jump target pointing at non-first instruction of the list. This jump target is used during jump
	 * instruction modification.
	 * @author Mikhail Ermakov
	 *
	 */
	static class TargetedInstructionList {

		public InstructionList il;
		public InstructionHandle jumpTarget;
		public boolean removeInstruction;

		public TargetedInstructionList(InstructionList il, InstructionHandle jumpTarget, boolean removeInstruction) {
			this.il = il;
			this.jumpTarget = jumpTarget;
			this.removeInstruction = removeInstruction;
		}

		public TargetedInstructionList(InstructionList il, InstructionHandle jumpTarget) {
			this(il, jumpTarget, false);
		}
	}

	public Mokka(String className, Set<String> methods) throws ClassNotFoundException
	{
		Logger.debug("Mokka initialization with class name " + className + " and methods " + methods);

		this.className = className;
		this.methods = methods;

		originCode = Repository.lookupClass(className);
		modifiedCode = new ClassGen(originCode);


	}

	/**
	 * Returns names of target class parent classes.
	 * @return set of parent classes names
	 */
	public Set<String> getParentClasses() {

		JavaClass parentClass = originCode;
		Set<String> result = new HashSet<String>();
		try {
			while ((parentClass = parentClass.getSuperClass()) != null) {
				result.add(parentClass.getClassName().replace('.', '/'));
			}
		} catch (ClassNotFoundException e) {
		}

		return result;
	}

	private List<Method> getTargetMethods() {

		List<Method> result = new ArrayList<Method>();
		for (Method method : originCode.getMethods()) {
			String name = method.getName();
			if (method.isNative()) {
				Logger.warn("method " + className + "." + name + " is native, skipping instrumentation");
				continue;
			}
			if (method.isAbstract())
			{
				Logger.warn("method " + className + "." + name + " is abstract, skipping instrumentation");
				continue;
			}
			if (!methods.contains(method.getName() + method.getSignature())) {
				Logger.warn("method " + className + "." + name + " is not included in method list, skipping instrumentation");
				continue;
			}
			if (method.getCode() == null) {
				Logger.warn("method " + className + "." + name + " is empty, skipping instrumentation");
			}
			result.add(method);
		}

		return result;
	}

	/**
	 * Instruments instance class file.
	 */
	public void instrument() throws ClassNotFoundException, IOException
	{
		if (Utility.isInstrumented(originCode)) {
			return;
		}

		List<CodeExceptionGen> exceptionHandlers;

		InstructionHandle ih, tempIH;
		ConstantPoolGen cpg = new ConstantPoolGen(originCode.getConstantPool());

		List<Method> methods = getTargetMethods();
		Milk milk = Milk.createMilk(modifiedCode, cpg);

		for (Method method : methods)
		{
			CoffeeMethodGen newMethod = new CoffeeMethodGen(method, className, cpg);
			LineNumberTable lnt = newMethod.getLineNumberTable(cpg);
			InstructionList il = newMethod.getInstructionList();
			List<InstructionHandle> instructionsToDelete = new ArrayList<InstructionHandle>();

			modifiedCode.removeMethod(method);
			milk.setCurrentMethod(newMethod);
			exceptionHandlers = Arrays.asList(newMethod.getExceptionHandlers());

			int line = -1;
			int previousLine = -2;

			for (ih = il.getStart(); ih != null;)
			{
				if (lnt != null) {
					line = lnt.getSourceLine(ih.getPosition());
					if (line != previousLine) {
						CoffeeMachine.codePositions.println("#PC " + Long.toHexString(pc) + " java " +
															className + "." + method.getName() + " " +
															originCode.getFileName() + " " + line);
						Milk.pc = pc ++;
						previousLine = line;
					} else {
						Milk.pc = pc - 1;
					}
				}

				if (Utility.isExceptionHandler(exceptionHandlers, ih)) {
					InstructionList exceptionHandlerList = milk.createExceptionHandlerList();

					if (exceptionHandlerList != null)
					{
						InstructionHandle target = il.insert(ih, exceptionHandlerList);
						for (CodeExceptionGen exceptionHandler : exceptionHandlers) {
							if (exceptionHandler.getHandlerPC() == ih) {
								exceptionHandler.setHandlerPC(target);
							}
						}
					}
				}
				TargetedInstructionList beforeInstructions = milk.createBeforeList(ih);
				TargetedInstructionList afterInstructions = milk.createAfterList(ih);

				il.insert(ih, beforeInstructions.il);

				if (beforeInstructions.jumpTarget != null)
				{
					for (BranchInstruction jumpTargeter : Utility.getJumpTargeters(ih))
					{
						jumpTargeter.setTarget(beforeInstructions.jumpTarget);
					}
				}

				tempIH = ih;
				ih = ih.getNext();

				if (afterInstructions != null)
				{
					il.append(tempIH, afterInstructions.il);
				}
				if (beforeInstructions.removeInstruction) {
					for (CodeExceptionGen exceptionHandler : newMethod.getExceptionHandlers()) {
						Utility.redirectExceptionHandler(exceptionHandler, tempIH, beforeInstructions.jumpTarget);
					}
					instructionsToDelete.add(tempIH);
				}
			}

			il.setPositions();

			for (InstructionHandle ihDel : instructionsToDelete) {
				try {
					il.delete(ihDel);
				} catch (TargetLostException e) {
				}
			}

			newMethod.fixLongJumps();

			modifiedCode.setConstantPool(cpg);
			newMethod.setMaxLocals();
			newMethod.setMaxStack();

			modifiedCode.addMethod(newMethod.getMethod());
		}

		double diff = ((double)Utility.getCodeSize(modifiedCode)) / Utility.getCodeSize(originCode);

		Logger.verbose(className + ".class:");
		Logger.verbose("    " + methods.size() + " method" + (methods.size() == 1 ? "" : "s") + " instrumented;");
		Logger.verbose("    size of code increased by " + (int)((diff - 1) * 100) + "%.");
	}

	/**
	 * Updates classfile in output location.
	 * @throws IOException
	 */
	public void updateClass() throws IOException {
		String path = Options.instrumentationDir + "/";

		modifiedCode.getJavaClass().dump(path + className + ".class");
	}

	/**
	 * Returns names of all methods in current class.
	 * @return names of all methods in current class
	 */
	public Set<String> getMethodNames()
	{
		Set<String> result = new HashSet<String>();

		for (Method method : modifiedCode.getMethods())
		{
			result.add(method.getName() + method.getSignature());
		}
		return result;
	}

	/**
	 * Modifies Daemon class to implement necessary output mechanism.
	 */
	public static void setDaemonOutput() {
		JavaClass daemonSource;
		ClassGen daemon = null;
		try {
			daemonSource = Repository.lookupClass("ru.ispras.coffeemachine.target.Daemon");
		} catch (ClassNotFoundException e) {
			Logger.error("cannot set Daemon output");
			return;
		}
		daemon = new ClassGen(daemonSource);
		ConstantPoolGen cpg = new ConstantPoolGen(daemonSource.getConstantPool());
		daemon.setConstantPool(cpg);
		CoffeeInstructionFactory ifc = CoffeeInstructionFactory.createFactory(daemon, cpg);

		MethodGen println = null;
		for (Method method : daemonSource.getMethods()) {
			if (method.getName().equals("println")) {
				println = new MethodGen(method, "Daemon", cpg);
				daemon.removeMethod(method);
				break;
			}
		}
		if (println == null) {
			return;
		}

		InstructionList il = println.getInstructionList();
		/* Daemon.println(String) is empty at this point */
		il.insert(ifc.createDaemonPrintln());
		println.setMaxLocals();
		println.setMaxStack();
		daemon.addMethod(println.getMethod());

		try {
			daemon.getJavaClass().dump(Options.instrumentationDir + "/ru/ispras/coffeemachine/target/Daemon.class");
		} catch (IOException e) {
			e.printStackTrace();
			Logger.error("cannot set Daemon output");
		}
	}

	/**
	 * Updates Concurrency output file name static field.
	 * @param fileName new file path value
	 */
	public static void setConcurrencyFile(String fileName)
	{
		JavaClass concSource;
		ClassGen conc = null;
		try {
			concSource = Repository.lookupClass("ru.ispras.coffeemachine.target.Concurrency");
		} catch (ClassNotFoundException e) {
			Logger.error("cannot set Concurrency output");
			return;
		}
		conc = new ClassGen(concSource);
		ConstantPoolGen cpg = new ConstantPoolGen(concSource.getConstantPool());
		conc.setConstantPool(cpg);

		MethodGen clinit = null;
		for (Method method : concSource.getMethods()) {
			if (method.getName().equals("<clinit>")) {
				clinit = new MethodGen(method, "Concurrency", cpg);
				conc.removeMethod(method);
				break;
			}
		}
		if (clinit == null) {
			return;
		}

		InstructionList il = clinit.getInstructionList();
		/* Daemon.println(String) is empty at this point */
		for (InstructionHandle ih : il.getInstructionHandles())
		{
			if ((ih.getInstruction().getOpcode() == Constants.ACONST_NULL) &&
					(ih.getNext().getInstruction().getOpcode() == Constants.INVOKESPECIAL))
			{
				ih.setInstruction(new PUSH(cpg, fileName).getInstruction());
			}
		}
		clinit.update();
		conc.addMethod(clinit.getMethod());

		try {
			conc.getJavaClass().dump(Options.instrumentationDir + "/ru/ispras/coffeemachine/target/Concurrency.class");
		} catch (IOException e) {
			e.printStackTrace();
			Logger.error("cannot set Concurrency output");
		}
	}

	public void setMethods(Set<String> methods)
	{
		this.methods = methods;
	}
}
